import cv2 as cv
import sys
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import numpy as np
import time
import datetime


def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


img_dir = '../../../../../large_data/pic/morphology'
img_name = 'i_bg_noise.png'
img_path1 = os.path.join(img_dir, img_name)
img_name = 'i_fg_noise.png'
img_path2 = os.path.join(img_dir, img_name)

spr = 3
spc = 4
spn = 0
plt.figure(figsize=[12, 6])

kernel = np.ones((5, 5), dtype=np.uint8)

sep('bg_noise')
img1 = cv.imread(img_path1, cv.IMREAD_COLOR)
print('bg_noise shape', img1.shape)
my_show_img(img1, 'bg_noise', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('opening')
opening = cv.morphologyEx(img1, cv.MORPH_OPEN, kernel)
my_show_img(opening, 'opening', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('erosion => dilation')
erosion = cv.erode(img1, kernel, iterations=1)
dilation = cv.dilate(erosion, kernel, iterations=1)
my_show_img(dilation, 'erosion => dilation', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('top hat')
tophat = cv.morphologyEx(img1, cv.MORPH_TOPHAT, kernel)
my_show_img(tophat, 'tophat', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('src - opening')
diff = cv.subtract(img1, opening)
my_show_img(diff, 'src - opening', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('fg_noise')
img2 = cv.imread(img_path2, cv.IMREAD_COLOR)
print('fg_noise shape', img2.shape)
my_show_img(img2, 'fg_noise', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('closing')
closing = cv.morphologyEx(img2, cv.MORPH_CLOSE, kernel)
my_show_img(closing, 'closing', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('dilation => erosion')
dilation = cv.dilate(img2, kernel, iterations=1)
erosion = cv.erode(dilation, kernel, iterations=1)
my_show_img(erosion, 'dilation => erosion', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('black hat')
blackhat = cv.morphologyEx(img2, cv.MORPH_BLACKHAT, kernel)
my_show_img(blackhat, 'blackhat', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('closing - src')
diff = cv.subtract(closing, img2)
my_show_img(diff, 'closing - src', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

plt.show()
